package com.aisino.utils;

import com.aisino.entity.TreeNode;
import com.baomidou.mybatisplus.core.toolkit.CollectionUtils;

import java.util.*;

public class TreeUtil {
    /**
     * 获取所有的根节点
     */
    public static <T extends TreeNode> List<T> getRootNode(List<T> treeNodeList, Set<Integer> nodeIdSet) {
        List<T> rootNodeList = new ArrayList<>();
        for (T treeNode : treeNodeList) {
            if (0 == treeNode.getParentId() || null == treeNode.getParentId()) {
                rootNodeList.add(treeNode);
                nodeIdSet.add(treeNode.getId());
            }
        }
        return rootNodeList;
    }

    /**
     * 构建树
     */
    public static <T extends TreeNode> T buildChildTree(T rootNode, List<T> treeNodeList, Set<Integer> nodeIdSet) {
        List<T> childNodeList = new ArrayList<>();
        for (T treeNode : treeNodeList) {
            if (!nodeIdSet.contains(treeNode.getId()) && treeNode.getParentId().equals(rootNode.getId())) {
                nodeIdSet.add(treeNode.getId());
                //这里相当于把当前节点又作为父节点 递归查询子节点
                childNodeList.add(buildChildTree(treeNode, treeNodeList, nodeIdSet));
            }
        }
        rootNode.setChildList((List<TreeNode>) childNodeList);
        return rootNode;
    }

    public static <T extends TreeNode> List<T> buildTree(List<T> treeNodeList) {
        if (CollectionUtils.isEmpty(treeNodeList)) {
            return Collections.emptyList();
        }
        //用来保存每一棵完整的树
        List<T> treeNodes = new ArrayList<T>();
        //用来保存已装填过的树节点ID 防止重复遍历
        Set<Integer> nodeIdSet = new HashSet<>();
        for (T root : getRootNode(treeNodeList, nodeIdSet)) {
            T treeNode = buildChildTree(root, treeNodeList, nodeIdSet);
            treeNodes.add(treeNode);
        }
        return treeNodes;
    }
}